# !/usr/bin/python3
# -*- coding: utf-8 -*-
""" 数据集分类
@Author         :  dax
@Version        :  
------------------------------------
@File           :  dataset_classify.py
@Description    :  
@CreateTime     :  2019/12/18 14:46
------------------------------------
@ModifyTime     :  
"""
import os
import shutil

import numpy as np

root_dir = r'C:\PycharmProjects\torchLearning'
target_dir = 'train'
typ = ['cat', 'dog']
data_dir = './data/'


def classify(begin, end):
    """ 数据集分类取样
    @param begin: 起始数
    @param end:  结束数
    """
    train_path = make_train_dir(basedir=root_dir, dataset_dir=target_dir)

    for tp in typ:
        tp_path = make_train_dir(train_path, tp)
        images = [tp + '.{}.jpg'.format(i) for i in np.arange(begin, end)]
        for img in images:
            src = os.path.join(data_dir, img)
            dst = os.path.join(tp_path, img)
            shutil.copyfile(src, dst)


def make_train_dir(basedir, dataset_dir):
    """ 创建 训练集文件夹

    @param basedir: root
    @param dataset_dir: dataset_dir
    """
    if not os.path.exists(basedir):
        os.mkdir(basedir)

    train_path = os.path.join(basedir, dataset_dir)

    if not os.path.exists(train_path):
        os.mkdir(train_path)
    return train_path


classify(0, 10)
